from Crypto.Util.number import getPrime, getRandomRange, inverse
from hashlib import sha256
import gmpy2


class PaillierZK:
    def __init__(self, key_length=512):
        # 每次实例化时都会生成新的密钥对
        self.generate_keys(key_length)

    def generate_keys(self, key_length):
        p = getPrime(key_length // 2)
        q = getPrime(key_length // 2)
        self.N = p * q
        self.N_sq = self.N ** 2
        self.lambda_ = gmpy2.lcm(p - 1, q - 1)
        self.g = self.N + 1
        self.mu = inverse(self.L(pow(self.g, self.lambda_, self.N_sq)), self.N)

    def L(self, x):
        return (x - 1) // self.N

    def encrypt_block(self, m: int):
        r = getRandomRange(1, self.N)
        c = pow(self.g, m, self.N_sq) * pow(r, self.N, self.N_sq) % self.N_sq

        a = getRandomRange(0, self.N)
        s = getRandomRange(1, self.N)
        u = pow(self.g, a, self.N_sq) * pow(s, self.N, self.N_sq) % self.N_sq

        e = int.from_bytes(sha256(str(c).encode() + str(u).encode()).digest(), 'big') % self.N
        z = (a + e * m) % self.N
        t = (s * pow(r, e, self.N)) % self.N

        return int(c), (int(u), int(e), int(z), int(t))

    def verify_block(self, c, proof):
        u, e, z, t = proof
        left = pow(self.g, z, self.N_sq) * pow(t, self.N, self.N_sq) % self.N_sq
        right = u * pow(c, e, self.N_sq) % self.N_sq
        return left == right

    def decrypt_block(self, c):
        return int((self.L(pow(c, self.lambda_, self.N_sq)) * self.mu) % self.N)
